import os
import torch
from torch.utils.data import Dataset, DataLoader
from PIL import Image
import numpy as np
from .augment_utils import get_train_augmentation, get_val_augmentation, get_inference_augmentation


def get_num_classes(config):
    """从配置中获取类别数"""
    if "name_classes" in config["data"]:
        return len(config["data"]["name_classes"])
    else:
        return config["data"]["num_classes"]


class SegmentationDataset(Dataset):
    """医学图像分割数据集"""
    
    def __init__(self, image_dir, label_dir, file_list, transform=None):
        """
        Args:
            image_dir (str): 图像文件夹路径
            label_dir (str): 标签文件夹路径
            file_list (str): 包含文件名的txt文件路径（文件名不带后缀）
            transform: 数据增强变换
        """
        self.image_dir = image_dir
        self.label_dir = label_dir
        self.transform = transform
        
        # 读取文件列表（文件名不带后缀）
        with open(file_list, 'r') as f:
            self.file_names = [line.strip() for line in f.readlines()]
    
    def __len__(self):
        return len(self.file_names)
    
    def __getitem__(self, idx):
        # 获取文件名（不带后缀）
        file_name = self.file_names[idx]
        
        # 构建完整的文件路径（添加.png后缀）
        image_path = os.path.join(self.image_dir, file_name + '.png')
        label_path = os.path.join(self.label_dir, file_name + '.png')
        
        # 加载图像
        image = Image.open(image_path).convert('RGB')
        
        # 加载标签（假设标签是灰度图像）
        label = Image.open(label_path).convert('L')
        
        # 应用数据增强
        if self.transform:
            image, label = self.transform(image, label)
        
        return image, label


def create_data_loaders(config, current_epoch=0):
    """
    创建训练、验证和测试数据加载器
    
    Args:
        config (dict): 配置字典
        current_epoch (int): 当前训练轮次，用于控制数据增强
    
    Returns:
        tuple: (train_loader, val_loader, test_loader)
    """
    
    # 获取数据增强变换
    train_transform = get_train_augmentation(config, current_epoch)
    val_transform = get_val_augmentation(config)
    test_transform = get_val_augmentation(config)  # 测试使用与验证相同的变换
    
    # 创建数据集
    train_dataset = SegmentationDataset(
        image_dir=config["data"]["image_dir"],
        label_dir=config["data"]["label_dir"],
        file_list=config["data"]["train_list"],
        transform=train_transform
    )
    
    val_dataset = SegmentationDataset(
        image_dir=config["data"]["image_dir"],
        label_dir=config["data"]["label_dir"],
        file_list=config["data"]["val_list"],
        transform=val_transform
    )
    
    test_dataset = SegmentationDataset(
        image_dir=config["data"]["image_dir"],
        label_dir=config["data"]["label_dir"],
        file_list=config["data"]["test_list"],
        transform=test_transform
    )
    
    # 创建数据加载器
    train_loader = DataLoader(
        train_dataset,
        batch_size=config["training"]["batch_size"],
        shuffle=True,
        num_workers=4,
        pin_memory=True
    )
    
    val_loader = DataLoader(
        val_dataset,
        batch_size=config["training"]["batch_size"],
        shuffle=False,
        num_workers=4,
        pin_memory=True
    )
    
    test_loader = DataLoader(
        test_dataset,
        batch_size=1, # config["training"]["batch_size"],
        shuffle=False,
        num_workers=4,
        pin_memory=True
    )
    
    return train_loader, val_loader, test_loader


def create_inference_loader(config, input_path=None, batch_size=1):
    """
    创建推理数据加载器
    
    Args:
        config (dict): 配置字典
        input_path (str): 输入路径，可以是图像目录或txt文件，如果为None则使用config中的路径
        batch_size (int): 批次大小
    
    Returns:
        DataLoader: 推理数据加载器
        list: 图像文件名列表
    """
    
    if input_path is None:
        input_path = config["inference"]["input_list"]
    
    # 获取推理变换
    inference_transform = get_inference_augmentation(config)
    
    # 判断输入路径是目录还是txt文件
    if os.path.isdir(input_path):
        # 如果是目录，获取图像文件列表（不带后缀）
        image_files = []
        for f in os.listdir(input_path):
            if f.lower().endswith((".png", ".jpg", ".jpeg", ".tiff", ".bmp")):
                # 去掉文件后缀
                file_name = os.path.splitext(f)[0]
                image_files.append(file_name)
        
        # 创建临时txt文件
        temp_list_path = os.path.join(input_path, "temp_inference_list.txt")
        with open(temp_list_path, "w") as f:
            for img_file in image_files:
                f.write(img_file + "\n")
        
        image_dir = input_path
        file_list = temp_list_path
        cleanup_temp = True
    else:
        # 如果是txt文件，直接使用
        image_dir = config["inference"]["image_dir"]
        file_list = input_path
        cleanup_temp = False
        
        # 读取文件名列表
        with open(file_list, "r") as f:
            image_files = [line.strip() for line in f.readlines()]
    
    # 创建推理数据集（标签目录设为图像目录，因为推理时不需要真实标签）
    inference_dataset = SegmentationDataset(
        image_dir=image_dir,
        label_dir=image_dir,  # 推理时标签目录设为图像目录
        file_list=file_list,
        transform=inference_transform
    )
    
    # 创建数据加载器
    inference_loader = DataLoader(
        inference_dataset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=2,
        pin_memory=True
    )
    
    # 清理临时文件
    if cleanup_temp:
        os.remove(temp_list_path)
    
    return inference_loader, image_files


class SingleImageDataset(Dataset):
    """单张图像数据集，用于单图推理"""
    
    def __init__(self, image_path, transform=None):
        self.image_path = image_path
        self.transform = transform
    
    def __len__(self):
        return 1
    
    def __getitem__(self, idx):
        image = Image.open(self.image_path).convert('RGB')
        
        if self.transform:
            # 为单图推理创建一个假的标签
            dummy_label = Image.new('L', image.size, 0)
            image, _ = self.transform(image, dummy_label)
        
        return image


def create_single_image_loader(config, image_path):
    """
    为单张图像创建数据加载器
    
    Args:
        config (dict): 配置字典
        image_path (str): 图像路径
    
    Returns:
        DataLoader: 单图数据加载器
    """
    
    inference_transform = get_inference_augmentation(config)
    
    dataset = SingleImageDataset(image_path, transform=inference_transform)
    
    loader = DataLoader(
        dataset,
        batch_size=1,
        shuffle=False,
        num_workers=1
    )
    
    return loader


if __name__ == "__main__":
    # 测试数据加载器
    import yaml
    
    # 创建一个假的配置
    fake_config = {
        "data": {
            "image_dir": "./data/images",
            "label_dir": "./data/labels",
            "train_list": "./data/train.txt",
            "val_list": "./data/val.txt",
            "test_list": "./data/test.txt",
            "image_size": [256, 256],
            "name_classes": ["Background", "Class1"]
        },
        "training": {
            "batch_size": 4
        },
        "augmentations": {
            "enabled": True,
            "augment_stop_epoch": 80,
            "flip": True,
            "rotate": 10,
            "scale": [0.8, 1.2]
        }
    }
    
    print("数据加载器配置完成")
    print("支持的功能:")
    print("1. 训练/验证/测试数据加载")
    print("2. 批量推理数据加载")
    print("3. 单图推理数据加载")
    print("4. 自动数据增强（根据epoch控制）")
    print("5. 文件名不带后缀的支持")
    print("6. 自动从name_classes获取类别数")

